import argparse
import os
import time
from pathlib import Path
import uuid
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords,
    xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
from utils.torch_utils import select_device, load_classifier, time_synchronized
import json

os.environ['KMP_DUPLICATE_LIB_OK']='True'

class yosDispatcher:



    def loadModel(self,weights,device):
        # Initialize
        set_logging()
        device = select_device(device)

        half = device.type != 'cpu'  # half precision only supported on CUDA

        # Load model
        model = attempt_load(weights, map_location=device)  # load FP32 model
        if half:
            model.half()  # to FP16

        return model,device,half


    def detect(save_img=False,model=None,device=None,half=False,ImgPath=None,
               outPath=None,imageSize=640,augment=None,conf_thres=0.4,iou_thres=0.5,
               classes=None,agnostic_nms=None,imgInfoMap=None):


        imgsz = check_img_size(imageSize, s=model.stride.max())  # check img_size

        # Set Dataloader
        vid_path, vid_writer = None, None

        save_img = True
        dataset = LoadImages(ImgPath, img_size=imgsz)

        # Get names and colors
        names = model.module.names if hasattr(model, 'module') else model.names
        colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

        # Run inference
        t0 = time.time()
        img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
        _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once

        detectionObjs=[]

        for path, img, im0s, vid_cap in dataset:
            img = torch.from_numpy(img).to(device)
            img = img.half() if half else img.float()  # uint8 to fp16/32
            img /= 255.0  # 0 - 255 to 0.0 - 1.0
            if img.ndimension() == 3:
                img = img.unsqueeze(0)

            # Inference
            t1 = time_synchronized()
            pred = model(img, augment=augment)[0]

            # Apply NMS
            pred = non_max_suppression(pred,conf_thres,iou_thres, classes=classes, agnostic=agnostic_nms)
            t2 = time_synchronized()


            # Process detections
            for i, det in enumerate(pred):  # detections per image

                p, s, im0 = path, '', im0s

                save_path = str(Path(outPath) / Path(p).name)
                s += '%gx%g ' % img.shape[2:]  # print string
                gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                if det is not None and len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class
                        s += '%g %ss, ' % (n, names[int(c)])  # add to string

                    # Write results
                    for *xyxy, conf, cls in reversed(det):

                        if save_img:  # Add bbox to image
                            label = '%s %.2f' % (names[int(cls)], conf)

                            labelTag=label.split(" ")

                            cordinateArray=[]

                            cordinateArray.append(int(xyxy[0]))
                            cordinateArray.append(int(xyxy[1]))
                            cordinateArray.append(int(xyxy[2]))
                            cordinateArray.append(int(xyxy[3]))

                            detectionItem={
                                "labelName":labelTag[0],
                                "score":labelTag[1],
                                "Coordinate":cordinateArray
                            }
                            detectionObjs.append(detectionItem)
                            plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)


                # Print time (inference + NMS)
                print('%sDone. (%.3fs)' % (s, t2 - t1))

                # Save results (image with detections)
                if save_img:
                    if dataset.mode == 'images':
                        cv2.imwrite(save_path, im0)
        print('Done. (%.3fs)' % (time.time() - t0))

        #组装返回数据
        resultData={
            "requestCode":uuid.uuid4(),
            "isDetectHurt":True if len(detectionObjs)>0 else False,
            "summary":s,
            "time":time.time(),
            "detectDuration":'%.3fS'%(t2 - t1),
            "detectionObjs":detectionObjs,
            'detectType':imgInfoMap['detectType'],
            'fileName':imgInfoMap['fileName'],
            'detectionImg':imgInfoMap['imgPrefix'] + "/" + imgInfoMap['processImgName']
        }
        return resultData


if __name__ == '__main__':


    yos=yosDispatcher()

    imgPath='../inference/images/zidane.jpg'
    outPath='../inference/output/'

    model,device,half=yos.loadModel('../weights/yolov5s.pt','cpu')
    with torch.no_grad():
          result=yos.detect(model=model,device=device,half=half,ImgPath=imgPath,outPath=outPath)
          print(result)
